import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.types
from IPython import display
class show():
    def __init__(self, legend, xlabel=None, ylabel=None, xlim=None, ylim=None, 
                 figsize=(4.5,3.5)) -> None: # 尺寸
        # self.legend = legend
        self.xlabel = xlabel
        self.ylabel = ylabel
        self.xlim = xlim
        self.ylim = ylim
        self.data = np.zeros(shape = (1,len(legend))) # 存储每个回合的数据，列是时间，行是不同的数据
        # self.data = [[0]*len(legend)]
        self.fmt = ['-', '--', '-.', ':']
        self.color = ["#66ccff","#FFAEC9","#26C755"]
        self.fig, self.axes = plt.subplots(1, 1, figsize=figsize)
        self.axes.set_xlabel(xlabel)
        self.axes.set_ylabel(ylabel)
        if self.xlim!=None:
            self.axes.set_xlim(xlim)
        if self.ylim!=None:
            self.axes.set_ylim(ylim)
        self.axes.grid()
        self.len = 1
        self.legend = legend
    def add(self, args):
        self.len+=1
        self.axes.legend(self.legend)
        if len(args)!=len(self.legend):
            assert(f"lenth of args {len(args)} no match length of legend{len(args)}")
        data:list = self.data.tolist()
        if isinstance(args[0],torch.Tensor):
            data.append([arg.to('cpu') for arg in args])
        else:
            data.append(args)
        self.data = np.array(data)
        length = len(self.legend)
        for i in range(length):
            self.axes.plot(list(range(1,self.len)),self.data[1:,i],color=self.color[i%3],linestyle=self.fmt[i//3],linewidth=1)
        display.clear_output(wait=True)
        display.display(self.fig)
    
    # def show_train(self):
    #     plt.cla()
    #     if self.xlim!=None:
    #         plt.xlim(self.xlim)
    #     if self.ylim!=None:
    #         plt.ylim(self.ylim)
    #     plt.xlabel(self.xlabel)
    #     plt.ylabel(self.ylabel)
    #     plt.grid()
    #     length = len(self.legend)
    #     for i in range(length):
    #         plt.plot(np.arange(len(self.data))[1:],self.data[1:,i],color = self.color[i],label = i+1)
    #         # plt.plot([1,2,3],[1,2,3],color = self.color[i],label = i+1)
    #     plt.legend(self.legend)
    #     plt.show()

    def eval_loss(self,loss,net,iter_data,showrate,device):
        losssum = 0.0
        num = 0
        net.to(device)
        for x, y in iter_data:
            x,y = x.to(device),y.to(device)
            y_h = net(x)
            l = loss(y_h, y)/len(y)
            losssum+=l
            num+=1
        l = losssum/num
        print(f"loss:{l}")
        return l*showrate
    
    def eval_acc(self,net,iter_data,device):

        right_sum = 0
        sums = 0
        net.to(device)
        for x, y in iter_data:
            x,y = x.to(device),y.to(device)
            y_h = net(x)
            ans = y_h.argmax(dim=1)
            maps = (ans==y)
            # print(ans)
            # print(y)

            right_sum += maps.type(dtype = torch.float32).sum()
            sums += len(y)
        acc = right_sum/sums
        # print(l)
        return acc

    def clear(self):
        self.data = np.zeros(shape = (1,len(self.legend)))



# my = show([1,2,3,4,5,6,7],"x","y")
# plt.ion()
# for i in range(10):
#     a = [1+i,2+i,3+i,4+i,5+i,6+i,7+i]
#     my.add(args=a)
#     print(np.arange(len(my.data))[1:])
#     print(my.data[1:,2])
#     plt.pause(0.5)